"""
消除模块：药物消除过程建模

该模块实现了药物在体内的消除过程，包括：
- 肝脏代谢清除
- 肾脏排泄清除
- 胆汁排泄清除
- 其他消除途径
- 消除动力学参数计算

作者：BlackCat@CPPO
版本：1.0.0
"""

import numpy as np

class EliminationModel:
    """
    药物消除模型类
    
    该类模拟药物在体内的消除过程，包括代谢和排泄途径。
    """
    
    def __init__(self, parameters):
        """
        初始化消除模型
        
        参数:
            parameters (dict): 模型参数，包含：
                - CL_hepatic: 肝脏清除率 (L/h)
                - CL_renal: 肾脏清除率 (L/h)
                - CL_biliary: 胆汁清除率 (L/h)
                - CL_other: 其他清除率 (L/h)
                - fu: 游离分数
                - BP: 血液-血浆浓度比
                - liver_flow: 肝脏血流量 (L/h)
                - kidney_flow: 肾脏血流量 (L/h)
        """
        self.params = parameters
        self.validate_parameters()
    
    def validate_parameters(self):
        """验证模型参数的有效性"""
        required_params = ['fu', 'BP']
        for param in required_params:
            if param not in self.params:
                raise ValueError(f"缺少必需参数：{param}")
    
    def calculate_elimination(self, individual_params, time_points, distribution_result):
        """
        计算药物消除过程
        
        参数:
            individual_params (dict): 个体特定参数
            time_points (np.array): 时间点数组
            distribution_result (dict): 分布阶段结果
            
        返回:
            dict: 消除过程结果
        """
        # 合并个体参数
        params = {**self.params, **individual_params}
        
        # 计算清除率
        clearance_rates = self.calculate_clearance_rates(params)
        
        # 计算消除半衰期，使用分布结果中的分布容积V_d
        elimination_half_life = self.calculate_elimination_half_life(params, clearance_rates, distribution_result)
        
        # 计算血浆浓度-时间曲线
        plasma_concentration = self.calculate_plasma_concentration(
            params, time_points, distribution_result, clearance_rates
        )
        
        # 计算组织消除
        tissue_elimination = self.calculate_tissue_elimination(
            params, time_points, distribution_result, clearance_rates
        )
        
        # 计算代谢产物
        metabolite_formation = self.calculate_metabolite_formation(
            params, time_points, distribution_result, clearance_rates
        )
        
        return {
            'clearance_rates': clearance_rates,
            'elimination_half_life': elimination_half_life,
            'plasma_concentration': plasma_concentration,
            'tissue_elimination': tissue_elimination,
            'metabolite_formation': metabolite_formation
        }
    
    def calculate_clearance_rates(self, params):
        """
        计算各种清除途径的清除率
        
        参数:
            params (dict): 模型参数
            
        返回:
            dict: 各清除途径的清除率
        """
        # 基础清除率
        CL_hepatic = params.get('CL_hepatic', 0.5)  # L/h
        CL_renal = params.get('CL_renal', 0.3)     # L/h
        CL_biliary = params.get('CL_biliary', 0.1) # L/h
        CL_other = params.get('CL_other', 0.05)    # L/h
        
        # 考虑血浆蛋白结合的影响
        fu = params['fu']
        unbound_fraction = max(fu, 0.01)  # 防止除零
        
        # 计算有效清除率
        hepatic_clearance = CL_hepatic * unbound_fraction
        renal_clearance = CL_renal * unbound_fraction
        biliary_clearance = CL_biliary * unbound_fraction
        other_clearance = CL_other * unbound_fraction
        
        return {
            'hepatic': hepatic_clearance,
            'renal': renal_clearance,
            'biliary': biliary_clearance,
            'other': other_clearance,
            'total': hepatic_clearance + renal_clearance + biliary_clearance + other_clearance
        }
    
    def calculate_elimination_half_life(self, params, clearance_rates, distribution_result=None):
        """
        计算消除半衰期
        
        参数:
            params (dict): 模型参数
            clearance_rates (dict): 清除率结果
            distribution_result (dict, optional): 分布阶段结果，用于获取计算得到的分布容积
            
        返回:
            dict: 消除半衰期信息
        """
        # 优先使用分布结果中的分布容积
        V_d = None
        if distribution_result is not None and isinstance(distribution_result, dict):
            # 从分布结果中获取V_d
            if 'distribution_volume' in distribution_result:
                V_d = distribution_result['distribution_volume']
            elif 'volume' in distribution_result:
                V_d = distribution_result['volume']
            elif 'Vss' in distribution_result:
                V_d = distribution_result['Vss']
            elif 'Vd' in distribution_result:
                V_d = distribution_result['Vd']
            elif 'parameters' in distribution_result and isinstance(distribution_result['parameters'], dict):
                # 检查分布结果中的参数
                dist_params = distribution_result['parameters']
                V_d = dist_params.get('distribution_volume') or dist_params.get('volume') or dist_params.get('Vss') or dist_params.get('Vd')
        
        # 验证V_d的有效性
        if V_d is None or not isinstance(V_d, (int, float)) or V_d <= 0:
            V_d = params.get('V_d', 15.0)  # 使用默认值
        
        # 确保V_d为有效正数
        V_d = max(float(V_d), 1e-6)  # L
        
        # 计算消除常数
        total_clearance = clearance_rates['total']
        elimination_constant = total_clearance / V_d
        
        # 计算半衰期
        t_half = 0.693 / max(elimination_constant, 1e-6)
        
        # 各途径的半衰期
        hepatic_t_half = 0.693 / max(clearance_rates['hepatic'] / V_d, 1e-6)
        renal_t_half = 0.693 / max(clearance_rates['renal'] / V_d, 1e-6)
        biliary_t_half = 0.693 / max(clearance_rates['biliary'] / V_d, 1e-6)
        other_t_half = 0.693 / max(clearance_rates['other'] / V_d, 1e-6)
        
        return {
            'total': t_half,
            'hepatic': hepatic_t_half,
            'renal': renal_t_half,
            'biliary': biliary_t_half,
            'other': other_t_half,
            'elimination_constant': elimination_constant,
            'V_d': V_d  # 返回实际使用的分布容积
        }
    
    def calculate_plasma_concentration(self, params, time_points, distribution_result, clearance_rates):
        """
        计算血浆浓度-时间曲线
        
        参数:
            params (dict): 模型参数
            time_points (np.array): 时间点数组
            distribution_result (dict): 分布阶段结果
            clearance_rates (dict): 清除率结果
            
        返回:
            np.array: 血浆浓度时间曲线
        """
        # 验证输入参数
        if not isinstance(time_points, np.ndarray):
            time_points = np.array(time_points)
        
        # 直接使用分布模块计算的血浆浓度，这些已经包含了吸收、分布和消除过程
        # 分布模块已经使用了一室模型公式：C(t) = (F * D * ka) / (V * (ka - ke)) * (e^(-ke*t) - e^(-ka*t))
        if isinstance(distribution_result, dict) and 'plasma_concentration' in distribution_result:
            plasma_concentration = distribution_result['plasma_concentration']
            
            # 验证浓度数据
            if not isinstance(plasma_concentration, np.ndarray):
                plasma_concentration = np.array(plasma_concentration)
            
            # 确保浓度非负
            plasma_concentration = np.maximum(plasma_concentration, 0.0)
            
            return plasma_concentration
        else:
            # 如果没有分布结果，使用简化的计算
            dose = params.get('dose', 100.0)
            V_d = params.get('V_d', 15.0)
            total_clearance = clearance_rates.get('total', 0.9)
            
            # 简化的单室模型
            ka = params.get('ka', 0.5)
            ke = total_clearance / V_d
            
            # 确保ka和ke不相等
            if abs(ka - ke) < 1e-6:
                ke = ke * 1.01
            
            # 计算生物利用度
            F = params.get('fa', 0.8) * params.get('fg', 0.9) * params.get('fh', 0.7)
            
            # 计算浓度
            plasma_concentration = (F * dose * ka / (V_d * (ka - ke))) * (
                np.exp(-ke * time_points) - np.exp(-ka * time_points)
            )
            
            return plasma_concentration
    
    def _get_V_d_from_distribution(self, distribution_result, params):
        """
        从分布结果中获取V_d，或使用默认值
        
        参数:
            distribution_result (dict): 分布阶段结果
            params (dict): 模型参数
            
        返回:
            float: 分布容积 (L)
        """
        V_d = None
        
        if isinstance(distribution_result, dict):
            # 从分布结果中获取V_d - 与distribution_module.py保持一致
            if 'distribution_volume' in distribution_result:
                V_d = distribution_result['distribution_volume']
            elif 'volume' in distribution_result:
                V_d = distribution_result['volume']
            elif 'Vss' in distribution_result:
                V_d = distribution_result['Vss']
            elif 'Vd' in distribution_result:
                V_d = distribution_result['Vd']
            elif 'parameters' in distribution_result and isinstance(distribution_result['parameters'], dict):
                dist_params = distribution_result['parameters']
                V_d = (dist_params.get('distribution_volume') or 
                      dist_params.get('volume') or 
                      dist_params.get('Vss') or 
                      dist_params.get('Vd'))
        
        # 验证V_d有效性
        if V_d is None or not isinstance(V_d, (int, float)) or V_d <= 0:
            V_d = params.get('Vss', 15.0)
        
        return max(float(V_d), 1e-6)
    
    def calculate_tissue_elimination(self, params, time_points, distribution_result, clearance_rates):
        """
        计算组织消除
        
        参数:
            params (dict): 模型参数
            time_points (np.array): 时间点数组
            distribution_result (dict): 分布阶段结果
            clearance_rates (dict): 清除率结果
            
        返回:
            dict: 组织消除结果
        """
        # 验证输入参数
        if not isinstance(time_points, np.ndarray):
            time_points = np.array(time_points)

        tissue_concentrations = distribution_result['tissue_concentrations']
        
        # 使用分布结果中的V_d计算消除常数
        V_d = self._get_V_d_from_distribution(distribution_result, params)
        total_clearance = max(clearance_rates['total'], 1e-6)
        elimination_constant = total_clearance / V_d
        elimination_constant = max(elimination_constant, 1e-6)
        
        tissue_elimination = {}
        
        for tissue, concentration in tissue_concentrations.items():
            if concentration is not None and len(concentration) > 0:
                # 验证浓度数据
                if not isinstance(concentration, np.ndarray):
                    concentration = np.array(concentration)
                
                # 使用分布相结束时的浓度作为消除相的起始浓度
                c0 = float(concentration[-1])
                
                # 处理无效值
                if np.isnan(c0) or np.isinf(c0) or c0 <= 0:
                    c0 = float(np.max(concentration))
                    if np.isnan(c0) or np.isinf(c0) or c0 <= 0:
                        c0 = 1.0
                c0 = max(c0, 1e-6)
                
                # 计算组织指数衰减
                tissue_elimination[tissue] = c0 * np.exp(-elimination_constant * time_points)
            else:
                # 使用默认的简单模型
                tissue_elimination[tissue] = np.exp(-elimination_constant * time_points)
            
            # 确保结果有效
            tissue_elimination[tissue] = np.maximum(tissue_elimination[tissue], 0.0)
        
        return tissue_elimination
    
    def calculate_metabolite_formation(self, params, time_points, distribution_result, clearance_rates):
        """
        计算代谢产物形成
        
        参数:
            params (dict): 模型参数
            time_points (np.array): 时间点数组
            distribution_result (dict): 分布阶段结果
            clearance_rates (dict): 清除率结果
            
        返回:
            dict: 代谢产物形成结果
        """
        # 验证输入参数
        if not isinstance(time_points, np.ndarray):
            time_points = np.array(time_points)

        # 获取吸收量 - 从分布结果中获取血浆浓度反向计算
        if isinstance(distribution_result, dict) and 'plasma_concentration' in distribution_result:
            plasma_concentration = distribution_result['plasma_concentration']
            
            # 验证浓度数据
            if not isinstance(plasma_concentration, np.ndarray):
                plasma_concentration = np.array(plasma_concentration)

            # 使用分布相结束时的浓度来计算有效剂量
            if len(plasma_concentration) > 0:
                # 使用分布相最后一个时间点的浓度
                c0 = float(plasma_concentration[-1])
                
                # 处理无效值
                if np.isnan(c0) or np.isinf(c0) or c0 <= 0:
                    c0 = float(np.max(plasma_concentration))
                    if np.isnan(c0) or np.isinf(c0) or c0 <= 0:
                        c0 = 0.929
                c0 = max(c0, 1e-6)
                
                V_d = self._get_V_d_from_distribution(distribution_result, params)
                effective_dose = c0 * V_d
            else:
                effective_dose = max(params.get('dose', 100.0), 1e-6)
        else:
            # 使用默认剂量
            effective_dose = max(params.get('dose', 100.0), 1e-6)

        # 计算代谢产物形成
        hepatic_clearance = max(clearance_rates['hepatic'], 1e-6)
        total_clearance = max(clearance_rates['total'], 1e-6)
        V_d = self._get_V_d_from_distribution(distribution_result, params)
        V_met = max(params.get('V_metabolite', 5.0), 1e-6)  # L

        # 计算代谢比例
        metabolic_fraction = min(hepatic_clearance / total_clearance, 1.0)

        # 使用梯形法计算累积代谢产物
        dt = time_points[1] - time_points[0] if len(time_points) > 1 else 0.1
        metabolite_formation_rate = metabolic_fraction * total_clearance * (effective_dose / V_d) * np.exp(-total_clearance * time_points / V_d)

        # 累积代谢产物
        cumulative_metabolite = np.cumsum(metabolite_formation_rate) * dt

        # 代谢产物浓度
        metabolite_concentration = cumulative_metabolite / V_met

        # 确保结果有效
        metabolite_concentration = np.maximum(metabolite_concentration, 0.0)
        cumulative_metabolite = np.maximum(cumulative_metabolite, 0.0)
        metabolite_formation_rate = np.maximum(metabolite_formation_rate, 0.0)

        return {
            'metabolite_concentration': metabolite_concentration,
            'cumulative_metabolite': cumulative_metabolite,
            'metabolic_fraction': metabolic_fraction,
            'formation_rate': metabolite_formation_rate
        }
    
    def calculate_hepatic_clearance(self, params):
        """
        计算肝脏清除率
        
        参数:
            params (dict): 模型参数
            
        返回:
            dict: 肝脏清除结果
        """
        # 肝脏生理学参数
        liver_weight = params.get('liver_weight', 1.8)  # kg
        liver_blood_flow = params.get('liver_flow', 1.35)  # L/h
        microsomal_protein = params.get('microsomal_protein', 45.0)  # mg/g liver
        
        # 药物特异性参数
        Clint = params.get('Clint', 10.0)  # 内在清除率 (L/h/kg microsomal protein)
        fu = params['fu']
        
        # 计算肝脏内在清除率
        intrinsic_clearance = Clint * microsomal_protein * liver_weight / 1000  # L/h
        
        # 计算肝脏清除率（well-stirred模型）
        hepatic_clearance = (liver_blood_flow * fu * intrinsic_clearance) / (
            liver_blood_flow + fu * intrinsic_clearance)
        
        # 计算肝脏提取率
        extraction_ratio = hepatic_clearance / liver_blood_flow
        
        return {
            'intrinsic_clearance': intrinsic_clearance,
            'hepatic_clearance': hepatic_clearance,
            'extraction_ratio': extraction_ratio,
            'liver_blood_flow': liver_blood_flow
        }
    
    def calculate_renal_clearance(self, params):
        """
        计算肾脏清除率
        
        参数:
            params (dict): 模型参数
            
        返回:
            dict: 肾脏清除结果
        """
        # 肾脏生理学参数
        kidney_weight = params.get('kidney_weight', 0.31)  # kg
        kidney_blood_flow = params.get('kidney_flow', 1.24)  # L/h
        # GFR = params.get('GFR', 7.5)  # L/h (成人)

        # 计算GFR（默认基于kidney_weight标准化）
        reference_weight = 0.3  # 标准肾脏重量 (kg)
        standard_GFR = 7.5  # 标准GFR (L/h)
        GFR = params.get('GFR', (kidney_weight / reference_weight) * standard_GFR)
        
        # 药物特异性参数
        fu = params['fu']
        renal_secretory_clearance = params.get('renal_secretory_clearance', 0.0)  # L/h
        
        # 计算肾小球滤过清除率
        glomerular_clearance = GFR * fu
        
        # 计算总肾脏清除率
        total_renal_clearance = glomerular_clearance + renal_secretory_clearance
        
        # 计算肾脏提取率
        renal_extraction_ratio = total_renal_clearance / kidney_blood_flow
        
        return {
            'glomerular_clearance': glomerular_clearance,
            'secretory_clearance': renal_secretory_clearance,
            'total_renal_clearance': total_renal_clearance,
            'extraction_ratio': renal_extraction_ratio,
            'kidney_blood_flow': kidney_blood_flow,
            'GFR': GFR
        }
    
    def calculate_biliary_clearance(self, params):
        """
        计算胆汁清除率
        
        参数:
            params (dict): 模型参数
            
        返回:
            dict: 胆汁清除结果
        """
        # 胆汁生理学参数
        bile_flow = params.get('bile_flow', 0.6)  # L/h
        
        # 药物特异性参数
        fu = params['fu']
        biliary_extraction_ratio = params.get('biliary_extraction_ratio', 0.1)
        
        # 计算胆汁清除率
        biliary_clearance = bile_flow * fu * biliary_extraction_ratio
        
        return {
            'biliary_clearance': biliary_clearance,
            'bile_flow': bile_flow,
            'extraction_ratio': biliary_extraction_ratio
        }
    
    def calculate_clearance_from_in_vitro_data(self, in_vitro_data):
        """
        根据体外数据计算清除率
        
        参数:
            in_vitro_data (dict): 体外实验数据
            
        返回:
            dict: 基于体外数据的清除率
        """
        # 微体实验数据
        if 'microsomal_data' in in_vitro_data:
            clint_microsomal = self._calculate_clint_from_microsomes(in_vitro_data['microsomal_data'])
        else:
            clint_microsomal = 0
        
        # 肝细胞实验数据
        if 'hepatocyte_data' in in_vitro_data:
            clint_hepatocytes = self._calculate_clint_from_hepatocytes(in_vitro_data['hepatocyte_data'])
        else:
            clint_hepatocytes = 0
        
        # 转运体实验数据
        if 'transporter_data' in in_vitro_data:
            transporter_clearance = self._calculate_transporter_clearance(in_vitro_data['transporter_data'])
        else:
            transporter_clearance = 0
        
        return {
            'microsomal_clearance': clint_microsomal,
            'hepatocyte_clearance': clint_hepatocytes,
            'transporter_clearance': transporter_clearance
        }
    
    def _calculate_clint_from_microsomes(self, microsomal_data):
        """根据微体数据计算内在清除率"""
        # 简化计算
        vmax = microsomal_data.get('Vmax', 0)
        km = microsomal_data.get('Km', 1e6)
        
        if km > 0:
            clint = vmax / km
        else:
            clint = 0
        
        return clint
    
    def _calculate_clint_from_hepatocytes(self, hepatocyte_data):
        """根据肝细胞数据计算内在清除率"""
        # 简化计算
        vmax = hepatocyte_data.get('Vmax', 0)
        km = hepatocyte_data.get('Km', 1e6)
        
        if km > 0:
            clint = vmax / km
        else:
            clint = 0
        
        return clint
    
    def _calculate_transporter_clearance(self, transporter_data):
        """计算转运体介导的清除率"""
        # 简化计算
        vmax = transporter_data.get('Vmax', 0)
        km = transporter_data.get('Km', 1e6)
        
        if km > 0:
            clint = vmax / km
        else:
            clint = 0
        
        return clint
    
    def calculate_elimination_kinetics(self, params, concentration_data):
        """
        计算消除动力学参数
        
        参数:
            params (dict): 模型参数
            concentration_data (np.array): 浓度数据
            
        返回:
            dict: 消除动力学参数
        """
        # 计算AUC
        time_points = np.linspace(0, len(concentration_data) * 0.1, len(concentration_data))
        AUC = np.trapezoid(concentration_data, time_points)
        
        # 计算清除率
        dose = params.get('dose', 100)
        clearance = dose / max(AUC, 1e-6)
        
        # 计算分布容积
        V_d = clearance * np.trapezoid(time_points * concentration_data, time_points) / max(AUC, 1e-6)
        
        # 计算半衰期
        # 使用最后几个点计算消除相半衰期
        last_points = min(5, len(concentration_data) // 2)
        if last_points > 1:
            x = time_points[-last_points:]
            y = np.log(concentration_data[-last_points:])
            
            # 线性回归
            A = np.vstack([x, np.ones(len(x))]).T
            slope, _ = np.linalg.lstsq(A, y, rcond=None)[0]
            
            t_half = 0.693 / max(abs(slope), 1e-6)
        else:
            t_half = np.nan
        
        return {
            'AUC': AUC,
            'clearance': clearance,
            'V_d': V_d,
            't_half': t_half,
            'mean_residence_time': V_d / max(clearance, 1e-6)
        }